import urllib.request  # 导入 urllib.request 模块，用于进行网络请求
import zipfile  # 导入 zipfile 模块，用于处理 ZIP 压缩文件
import os  # 导入 os 模块，用于操作系统相关的操作
from pathlib import Path  # 导入 Path 类，用于处理文件路径
import pandas as pd  # 导入 pandas 库，并使用别名 pd
import torch
from torch.utils.data import Dataset  # 从PyTorch中导入Dataset类
import tiktoken  # 导入tiktoken库
from torch.utils.data import DataLoader  # 从PyTorch中导入DataLoader类
from gpt2_download import download_and_load_gpt2
from gpt2_tools import GPTModel, load_weights_into_gpt, replace_linear_with_lora, plot_losses
import time
import matplotlib.pyplot as plt


def download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path):
    # 如果数据文件已经存在
    if data_file_path.exists():
        # 打印提示信息并跳过下载和解压步骤
        print(f"{data_file_path} already exists. Skipping download and extraction.")
        return

    # 使用 urllib.request.urlopen 打开 URL 进行网络请求
    with urllib.request.urlopen(url) as response:
        # 以二进制写入模式打开 ZIP 文件
        with open(zip_path, "wb") as out_file:
            # 将网络响应的内容写入 ZIP 文件
            out_file.write(response.read())
            # 这一步完成了文件的下载

    # 使用 zipfile.ZipFile 打开 ZIP 文件
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        # 将 ZIP 文件解压到指定的路径
        zip_ref.extractall(extracted_path)

    # 获取原始文件路径
    original_file_path = Path(extracted_path) / "SMsSpamCollection"
    # 将原始文件重命名为带有.tsv 扩展名的文件
    os.rename(original_file_path, data_file_path)

def create_balanced_dataset(df):
    # 计算标签为"spam"的实例数量
    num_spam = df[df["Label"] == "spam"].shape[0]
    # 从标签为"ham"的实例中随机采样，数量与"spam"的实例数量相同，设置随机种子为 123
    ham_subset = df[df["Label"] == "ham"].sample(num_spam, random_state=123)
    # 将标签为"spam"的子集和随机采样后的"ham"子集合并
    balanced_df = pd.concat([df[df["Label"] == "spam"], ham_subset])
    # 打印平衡后数据集中的标签数量分布
    print(balanced_df["Label"].value_counts())
    # 返回平衡后的数据集
    return balanced_df

def random_split(df, train_frac, validation_frac):
    # 随机打乱整个 DataFrame，设置随机种子为 123
    df = df.sample(frac=1, random_state=123).reset_index(drop=True)
    # 计算验证集的结束索引位置，通过数据集长度乘以验证集比例得到验证集长度，再加上训练集结束索引位置
    # 根据比例确定训练集结束索引位置
    train_end = int(len(df) * train_frac)
    # 按照索引位置分割 DataFrame，得到训练集
    train_df = df[:train_end]
    validation_end = train_end + int(len(df) * validation_frac)
    # 按照索引位置分割 DataFrame，得到验证集
    validation_df = df[train_end:validation_end]
    # 按照索引位置分割 DataFrame，得到测试集（剩余部分）
    test_df = df[validation_end:]
    # 返回训练集、验证集和测试集
    return train_df, validation_df, test_df

class SpamDataset(Dataset):
    # 初始化方法，接收 CSV 文件路径、标记器、最大长度和填充标记 ID 作为参数
    def __init__(self, csv_file, tokenizer, max_length=None, pad_token_id=50256):
        # 读取 CSV 文件数据并存入 self.data
        self.data = pd.read_csv(csv_file)
        # 对文本进行预标记化，将每个文本用标记器编码后存入 self.encoded_texts 列表
        self.encoded_texts = [tokenizer.encode(text) for text in self.data["Text"]]
        # 如果没有传入最大长度，则计算数据集中最长编码序列的长度并赋值给 self.max_length
        if max_length is None:
            self.max_length = self._longest_encoded_length()
        else:
            # 否则使用传入的最大长度
            self.max_length = max_length
        # 如果序列长度超过最大长度，则截断
        self.encoded_texts = [encoded_text[:self.max_length] if len(encoded_text) > self.max_length else encoded_text for encoded_text in self.encoded_texts]
        # 对长度不足最大长度的序列进行填充
        self.encoded_texts = [encoded_text + [pad_token_id] * (self.max_length - len(encoded_text)) for encoded_text in self.encoded_texts]

    # 根据索引获取数据集中的一个样本，返回编码后的文本和对应的标签
    def __getitem__(self, index):
        encoded = self.encoded_texts[index]
        label = self.data.iloc[index]["Label"]
        return (torch.tensor(encoded, dtype=torch.long), torch.tensor(label, dtype=torch.long))

    # 返回数据集的长度
    def __len__(self):
        return len(self.data)

    # 计算数据集中最长编码序列的长度的私有方法
    def _longest_encoded_length(self):
        max_length = 0
        for encoded_text in self.encoded_texts:
            encoded_length = len(encoded_text)
            if encoded_length > max_length:
                max_length = encoded_length
        return max_length

def calc_accuracy_loader(data_loader, model, device, num_batches=None):
    # 将模型设置为评估模式
    model.eval()
    correct_predictions, num_examples = 0, 0
    # 如果没有指定批次数量，则默认为数据加载器的长度
    if num_batches is None:
        num_batches = len(data_loader)
    else:
        # 遍历数据加载器中的批次
        for i, (input_batch, target_batch) in enumerate(data_loader):
            num_batches = min(num_batches, len(data_loader))
            if i < num_batches:
                # 将输入和目标批次移动到指定设备
                input_batch = input_batch.to(device)
                target_batch = target_batch.to(device)
                # 获取模型输出的最后一个输出标记的对数几率
                with torch.no_grad():
                    logits = model(input_batch)[:,-1,:]
                # 获取预测的标签
                predicted_labels = torch.argmax(logits, dim=-1)
                # 计算正确预测的数量和总样本数量
                num_examples += predicted_labels.shape[0]
                correct_predictions += (predicted_labels == target_batch).sum().item()
            else:
                break
    # 返回准确率
    return correct_predictions / num_examples

def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch = input_batch.to(device)
    target_batch = target_batch.to(device)
    logits = model(input_batch)[:,-1,:]
    loss = torch.nn.functional.cross_entropy(logits, target_batch)
    return loss

def calc_loss_loader(data_loader, model, device, num_batches=None):
    # 初始化总损失为 0
    total_loss = 0.
    # 如果数据加载器长度为 0，则返回 NaN
    if len(data_loader) == 0:
        return float("nan")
    elif num_batches is None:
        # 如果没有指定批次数量，则默认为数据加载器的长度
        num_batches = len(data_loader)
    else:
        # 确保批次数量不超过数据加载器中的批次数量
        num_batches = min(num_batches, len(data_loader))
        for i, (input_batch, target_batch) in enumerate(data_loader):
            if i < num_batches:
                # 计算当前批次的损失
                loss = calc_loss_batch(input_batch, target_batch, model, device)
                total_loss += loss.item()
            else:
                break
    # 返回平均损失
    return total_loss / num_batches

def evaluate_model(model, train_loader, val_loader, device, eval_iter):
    """
    评估模型函数。

    参数：
    - model：待评估的模型。
    - train_loader：训练数据加载器。
    - val_loader：验证数据加载器。
    - device：设备（如 CPU 或 GPU）。
    - eval_iter：评估迭代次数。

    返回：
    - train_loss：训练损失。
    - val_loss：验证损失。
    """
    model.eval()
    with torch.no_grad():
        train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
        val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
        model.train()
    return train_loss, val_loss

def train_classifier_simple(
        model, train_loader, val_loader, optimizer, device,
        num_epochs, eval_freq, eval_iter
):
    # 初始化用于记录损失和示例数量的列表
    train_losses, val_losses, train_accs, val_accs = [], [], [], []
    examples_seen, global_step = 0, -1

    # 主训练循环
    for epoch in range(num_epochs):
        model.train()  # 将模型设置为训练模式

        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()  # 重置上一个批次的梯度
            loss = calc_loss_batch(
                input_batch, target_batch, model, device  # 计算当前批次的损失
            )
            loss.backward()  # 计算损失的梯度
            optimizer.step()  # 使用梯度更新模型权重
            examples_seen += input_batch.shape[0]  # 更新已看到的示例数量
            global_step += 1

            # 可选的评估步骤
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device, eval_iter
                )
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                print(
                    f"Ep {epoch + 1} (Step {global_step:06d}): "
                    f"Train loss {train_loss:.3f}, "
                    f"Val loss {val_loss:.3f}"
                )

        # 每个epoch计算一次准确率
        train_accuracy = calc_accuracy_loader(
            train_loader, model, device, num_batches=eval_iter
        )
        val_accuracy = calc_accuracy_loader(
            val_loader, model, device, num_batches=eval_iter
        )

        print(f"Training accuracy: {train_accuracy * 100:.2f}% ", end="")
        print(f"Validation accuracy: {val_accuracy * 100:.2f}%")
        train_accs.append(train_accuracy)
        val_accs.append(val_accuracy)

    return train_losses, val_losses, train_accs, val_accs, examples_seen

def plot_values(
    epochs_seen, examples_seen, train_values, val_values,
    label="loss"
):
    # 创建图形和轴对象，设置图形大小
    fig, ax1 = plt.subplots(figsize=(5, 3))

    # 绘制训练和验证的损失或准确率随epoch变化的曲线
    ax1.plot(epochs_seen, train_values, label=f"Training {label}")
    ax1.plot(
        epochs_seen, val_values, linestyle="-.",
        label=f"Validation {label}"
    )
    ax1.set_xlabel("Epochs")  # 设置x轴标签为"Epochs"
    ax1.set_ylabel(label.capitalize())  # 设置y轴标签为参数label的首字母大写形式
    ax1.legend()  # 显示图例

    # 创建一个次坐标轴，用于显示"examples_seen"为x轴
    ax2 = ax1.twiny()
    ax2.plot(examples_seen, train_values, alpha=0)  # 添加不可见曲线对齐刻度
    ax2.set_xlabel("Examples seen")  # 设置第二个x轴标签为"Examples seen"

    fig.tight_layout()  # 调整布局以留出显示空间
    plt.savefig(f"{label}-plot.pdf")  # 将图形保存为PDF文件
    plt.show()  # 显示图形

# 数据集的URL
url = \
    "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"
# 压缩包路径
zip_path = "sms_spam_collection.zip"
# 解压后的文件夹路径
extracted_path = "sms_spam_collection"
# 数据文件路径
data_file_path = Path(extracted_path) / "SMSSpamCollection.tsv"

# 下载并解压数据集
download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)

# 读取数据文件，将其加载为DataFrame
df = pd.read_csv(
    data_file_path, sep="\t", header=None, names=["Label", "Text"]  # 指定列名
)

# 创建平衡的数据集（使ham和spam数量相等）
balanced_df = create_balanced_dataset(df)

# 将标签"ham"和"spam"映射为0和1
balanced_df["Label"] = balanced_df["Label"].map({"ham": 0, "spam": 1})

# 将数据集划分为训练集、验证集和测试集，比例为0.7:0.1:0.2
train_df, validation_df, test_df = random_split(balanced_df, 0.7, 0.1)

# 将训练集保存为CSV文件
train_df.to_csv("train.csv", index=None)

# 将验证集保存为CSV文件
validation_df.to_csv("validation.csv", index=None)

# 将测试集保存为CSV文件
test_df.to_csv("test.csv", index=None)

# 获取GPT-2编码器
tokenizer = tiktoken.get_encoding("gpt2")

# 创建训练集数据集
train_dataset = SpamDataset(
    "train.csv",  # 训练集CSV文件路径
    max_length=None,  # 最大长度未指定，保持原始长度
    tokenizer=tokenizer  # 使用的分词器
)

# 创建验证集数据集
val_dataset = SpamDataset(
    "validation.csv",  # 验证集CSV文件路径
    max_length=train_dataset.max_length,  # 与训练集一致的最大长度
    tokenizer=tokenizer  # 使用的分词器
)

# 创建测试集数据集
test_dataset = SpamDataset(
    "test.csv",  # 测试集CSV文件路径
    max_length=train_dataset.max_length,  # 与训练集一致的最大长度
    tokenizer=tokenizer  # 使用的分词器
)

# 设置DataLoader参数
num_workers = 0  # 设置工作线程数为0（在单机训练时通常为0）
batch_size = 8  # 每个batch的大小

# 设置随机种子，保证结果可复现
torch.manual_seed(123)

# 创建训练集的DataLoader
train_loader = DataLoader(
    dataset=train_dataset,  # 使用训练集数据集
    batch_size=batch_size,  # 每个batch的大小
    shuffle=True,  # 是否打乱数据
    num_workers=num_workers,  # 工作线程数
    drop_last=True  # 是否丢弃最后不足一个batch的数据
)

# 创建验证集的DataLoader
val_loader = DataLoader(
    dataset=val_dataset,  # 使用验证集数据集
    batch_size=batch_size,  # 每个batch的大小
    num_workers=num_workers,  # 工作线程数
    drop_last=False  # 不丢弃最后不足一个batch的数据
)

# 创建测试集的DataLoader
test_loader = DataLoader(
    dataset=test_dataset,  # 使用测试集数据集
    batch_size=batch_size,  # 每个batch的大小
    num_workers=num_workers,  # 工作线程数
    drop_last=False  # 不丢弃最后不足一个batch的数据
)

CHOOSE_MODEL ="gpt2-small (124M)"
INPUT_PROMPT="Every effort moves"

BASE_CONFIG={
    "vocab_size": 50257, "drop_rate": 0.0, "qkv_bias": True, "context_length": 1024
}
model_configs={
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads":12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-x1 (1558M)":{"emb_dim": 1600, "n_layers":48, "n_heads": 25}
}
BASE_CONFIG.update(model_configs[CHOOSE_MODEL])

model_size = CHOOSE_MODEL.split(" ")[-1].lstrip("(").rstrip(")")
settings, params = download_and_load_gpt2(model_size=model_size, models_dir="gpt2")

model = GPTModel(BASE_CONFIG)
load_weights_into_gpt(model, params)
model.eval()

torch.manual_seed(123)
num_classes = 2
model.out_head = torch.nn.Linear(in_features=BASE_CONFIG["emb_dim"], out_features=num_classes)
for param in model.parameters():
    param.requires_grad = False
replace_linear_with_lora(model, rank=16, alpha=16)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

start_time = time.time()
torch.manual_seed(123)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.1)
num_epochs = 5
train_losses, val_losses, train_accs, val_accs, examples_seen = \
train_classifier_simple(
    model, train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs, eval_freq=50, eval_iter=5
)
end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
examples_seen_tensor = torch.linspace(0, examples_seen, len(train_losses))
plot_values(
    epochs_tensor, examples_seen_tensor,
    train_losses, val_losses, label="loss"
)






